import akshare as ak
import pandas as pd
from retry import retry

from python_quant.stock.z_t_b import recent_trading_day_zt_before_1030


def today_amount_in_without_small_amount(in_rate):
    """
    主力资金、 超大单、大单、中单净买入比例大于in_rate
    """

    amount_df = ak.stock_individual_fund_flow_rank(indicator="今日")
    amount_clear_df = amount_df[
        (amount_df['今日主力净流入-净占比'].str.contains('-') == True)
        | (amount_df['今日超大单净流入-净占比'].str.contains('-') == True)
        | (amount_df['今日中单净流入-净占比'].str.contains('-') == True)
        | (amount_df['今日大单净流入-净占比'].str.contains('-') == True)
        ]

    set_diff_df = pd.concat([amount_df, amount_clear_df, amount_clear_df]).drop_duplicates(keep=False)

    set_diff_df = set_diff_df[
        (set_diff_df['今日中单净流入-净占比'].astype('float') > in_rate) &
        (set_diff_df['今日大单净流入-净占比'].astype('float') > in_rate) &
        (set_diff_df['今日主力净流入-净占比'].astype('float') > in_rate) & (
                set_diff_df['今日超大单净流入-净占比'].astype('float') > in_rate)]

    amount_in_codes_df = set_diff_df['代码']
    return amount_in_codes_df


def today_all_amount_in(in_rate):
    """
        筛选出实时主力资金净买入占比大于in_rate的个股
    """
    amount_df = ak.stock_individual_fund_flow_rank(indicator="今日")
    amount_df2 = amount_df[
        (amount_df['今日主力净流入-净占比'].str.contains('-') == True)
        & (amount_df['今日超大单净流入-净占比'].str.contains('-') == True)
        & (amount_df['今日超大单净流入-净占比'].str.contains('-') == True)
        & (amount_df['今日中单净流入-净占比'].str.contains('-') == True)
        & (amount_df['今日小单净流入-净占比'].str.contains('-') == True)
        ]
    set_diff_df = pd.concat([amount_df, amount_df2, amount_df2]).drop_duplicates(keep=False)
    set_diff_df = set_diff_df[
        (set_diff_df['今日小单净流入-净占比'].astype('float') > in_rate) &
        (set_diff_df['今日中单净流入-净占比'].astype('float') > in_rate) &
        (set_diff_df['今日大单净流入-净占比'].astype('float') > in_rate) &
        (set_diff_df['今日主力净流入-净占比'].astype('float') > in_rate) & (
                set_diff_df['今日超大单净流入-净占比'].astype('float') > in_rate)]
    return set_diff_df['代码']


def amount5_in(in_rate):
    """
        筛选出实时主力资金净买入占比大于in_rate的个股
    """
    amount_df = ak.stock_individual_fund_flow_rank(indicator="5日")
    amount_df2 = amount_df[
        (amount_df['5日主力净流入-净占比'].str.contains('-') == True) & (
                amount_df['5日超大单净流入-净占比'].str.contains('-') == True)]
    set_diff_df = pd.concat([amount_df, amount_df2, amount_df2]).drop_duplicates(keep=False)
    set_diff_df = set_diff_df[
        (set_diff_df['5日主力净流入-净占比'].astype('float') > in_rate) & (
                set_diff_df['5日超大单净流入-净占比'].astype('float') > in_rate)]
    return set_diff_df['代码']


def amount10_in(in_rate):
    """
        筛选出实时主力资金净买入占比大于10%的个股,股价小于10的个股
    """
    amount_df = ak.stock_individual_fund_flow_rank(indicator="10日")
    amount_df2 = amount_df[
        (amount_df['10日主力净流入-净占比'].str.contains('-') == True) & (
                amount_df['10日超大单净流入-净占比'].str.contains('-') == True)]
    set_diff_df = pd.concat([amount_df, amount_df2, amount_df2]).drop_duplicates(keep=False)
    set_diff_df = set_diff_df[
        (set_diff_df['10日主力净流入-净占比'].astype('float') > in_rate) & (
                set_diff_df['10日超大单净流入-净占比'].astype('float') > in_rate)]
    return set_diff_df


@retry(tries=10, delay=10, backoff=1)
def amount_stocks():
    """
    资金面整合
    """
    today_amount = today_amount_in_without_small_amount(0)
    zt_1030 = recent_trading_day_zt_before_1030()
    amount_in_codes = pd.merge(today_amount, zt_1030, on='代码', how='outer')
    return amount_in_codes
